package com.umbrella.robot.transformer;

import javassist.ClassPool;
import javassist.CtClass;
import javassist.CtMethod;
import javassist.Modifier;
import javassist.bytecode.CodeAttribute;
import javassist.bytecode.LocalVariableAttribute;
import javassist.bytecode.MethodInfo;

import java.io.ByteArrayInputStream;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.security.ProtectionDomain;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;

/**
 * @author Kim
 * @create 2018-01-30 18:02
 **/
public class WatchTransformer implements ClassFileTransformer {
    @Override
    public byte[] transform(ClassLoader loader, String className,
                            Class classBeingRedefined, ProtectionDomain protectionDomain,
                            byte[] classfileBuffer) throws IllegalClassFormatException {
        byte[] byteCode = classfileBuffer;

        // since this transformer will be called when all the classes are
        // loaded by the classloader, we are restricting the instrumentation
        // using if block only for the Lion class
        if ("com/rong/kim/common/Lion".equals(className)) {
            System.out.println("WatchTransformer Instrumenting......");
            try {
                ClassPool classPool = ClassPool.getDefault();
                CtClass ctClass = classPool.makeClass(new ByteArrayInputStream(
                    classfileBuffer));

                CtMethod[] methods = ctClass.getDeclaredMethods();
                for (CtMethod method : methods) {
                    String modifier = Modifier.toString(method.getModifiers());
                    CtClass returnType = method.getReturnType();
                    String name = method.getName();
                    CtClass[] parameterTypes = method.getParameterTypes();

                    //System.out.print("\n" + modifier + " " + returnType.getName() + " " + name + " (");

                    MethodInfo methodInfo = method.getMethodInfo();
                    CodeAttribute codeAttribute = methodInfo.getCodeAttribute();
                    LocalVariableAttribute attribute = (LocalVariableAttribute) codeAttribute.getAttribute(LocalVariableAttribute.tag);
                    int pos = Modifier.isStatic(method.getModifiers()) ? 0 : 1;
                    List<String> paras = new ArrayList<>();
                    for (int i = 0; i < method.getParameterTypes().length; i++) {
                        //System.out.print(parameterTypes[i] + " " + attribute.variableName(i + pos));
                        if (i < method.getParameterTypes().length - 1) {
                            //System.out.print(",");
                        }

                        paras.add(attribute.variableName(i + pos));
                    }

                    //System.out.print(")");

                    method.addLocalVariable("paras", classPool.get("java.lang.String"));
                    method.addLocalVariable("fw", classPool.get("java.io.FileWriter"));
                    method.addLocalVariable("bw", classPool.get("java.io.BufferedWriter"));
                    method.insertBefore("fw = new java.io.FileWriter(\"bw.txt\", true);" +
                        "bw = new java.io.BufferedWriter(fw);" +
                        "paras += \"\\n\";" +
                        "bw.write(paras);" +
                        "bw.flush();" +
                        "bw.close();");
                    Collections.reverse(paras);
                    for (String para : paras) {
                        method.insertBefore("System.out.println(\"" + para + ": \" + " + para + ");");
                        String msg = "paras += (\"" + para + "-> [\" +" + para + " + \"]  \");";
                        method.insertBefore(msg);
                    }

                    method.insertBefore("paras = \"\";");

                    CtClass[] exceptionTypes = method.getExceptionTypes();
                    if (exceptionTypes.length > 0) {
                        //System.out.print(" throws ");
                        int j = 0;
                        for (CtClass cl : exceptionTypes) {
                            //System.out.print(cl.getName());
                            if (j < exceptionTypes.length - 1) {
                                //System.out.print(",");
                            }
                            j++;
                        }
                    }
                }

                byteCode = ctClass.toBytecode();
                ctClass.detach();
                System.out.println("WatchTransformer Instrumentation complete.");
            } catch (Throwable ex) {
                System.out.println("Exception: " + ex);
                ex.printStackTrace();
            }
        }
        return byteCode;
    }
}
